import os
import tkinter as tk
from tkinter import *
from tkinter.filedialog import askopenfilename
import tkinter.font as tkFont
from tkinter.ttk import Separator
from PIL import Image, ImageTk
import torch
from networks.ClassicNetwork.VGGNet import VGG16
from networks.ClassicNetwork.AlexNet import AlexNet
from torchvision import transforms
# from model import AlexNet
from networks.ClassicNetwork.ResNet import ResNet50
font1 = {'family': 'Times New Roman',
         'weight': 'normal',
         'size': 18,
        }
font2 = {'family': 'Times New Roman',
         'weight' : 'normal',
         'size': 12,
        }

colors = ['orange', 'navy', 'red', 'blue']

class MainWindow:
    def __init__(self):
        self.root = tk.Tk()
        #self.root.title('阿尔兹海默症诊断该系统')
        self.root.overrideredirect(1)  #隐藏标题框
     #   canvas = tk.Canvas(self.root , height=450, width=600)
      #  imagefile = tk.PhotoImage(file='背景.png')
       # image = canvas.create_image(0, 0, anchor='nw', image=imagefile)
        #canvas.pack(side=tk.LEFT)
        #self.root.geometry('600x600')
        self.root.geometry("670x450+250+100")
        #self.frame_left = tk.Frame(self.root, padx=1, pady=5)
        self.frame_left = tk.Frame(self.root)
        self.frame_left.pack(fill='x',side='top')
        #self.frame_left.pack(padx=5, pady=10, fill="y", side=tk.LEFT)
        #canvas = tk.Canvas(self.frame_left, height=500, width=400)
        self.imagefile1 = tk.PhotoImage(file='背景.png')
        #image1 = canvas.create_image(0, 0, anchor='nw', image=imagefile1)
        #canvas.pack(side='top')
        self.canvas = tk.Canvas(self.frame_left, bg='gray',  height=450, width=600)
        self.canvas.pack(fill='x', expand='yes')
        self.canvas.create_image(0, 0, anchor='nw', image=self.imagefile1)  #背景图片
        ft = tkFont.Font(family='楷体', size=28, weight=tkFont.BOLD)
        self.canvas.create_text(40, 50, anchor='nw', font=ft, fill="white", text='欢迎来到阿尔兹海默症诊断系统')
        self.canvas.create_line(0, 120,
                                670, 120, fill='white')  #分割线
        self.canvas.create_line(350, 150,
                                350, 400, fill='gray')  #分割线
        self.canvas1 = tk.Canvas(self.canvas, bg='gray',  height=200, width=300)   #图片显示框画布
        self.canvas1.place(x=10, y=200)
        self.menubar = tk.Menu(self.root, activeborderwidth=20)   #菜单栏
        self.filemenu = tk.Menu(self.menubar, tearoff=0)
        self.menubar.add_cascade(label='选择模型',font=font1, menu=self.filemenu)
        self.menubar.add_command(label='退出', command=self.root.quit)
        self.filemenu .add_command(label='ResNet模型', command=self.loadModel1)
        self.filemenu.add_command(label='VGG模型', command=self.loadModel2)
        self.filemenu.add_command(label='AlexNet模型', command=self.loadModel3)
        self.filemenu.add_separator()  # 添加一条分隔线
        self.root.config(menu=self.menubar)
        self.button_loadModel = tk.Button(self.frame_left,activebackground='#AF7AC5',activeforeground='RED', font = font1, text='加载模型', command=self.loadModel)
        self.button_loadModel.place(x=10, y=150)
        #self.button_loadModel.pack(fill="x")
        self.button_loadModel.config(state=tk.DISABLED)

        self.button_loadImage = tk.Button(self.frame_left,activebackground='#AF7AC5',activeforeground='RED',  font = font1, text='加载图像', command=self.loadimg)
       # self.button_loadImage.pack(fill="x")
        self.button_loadImage .place(x=110, y=150)
        self.button_loadImage.config(state=tk.DISABLED)

        self.button_predict = tk.Button(self.frame_left,activebackground='#AF7AC5',activeforeground='RED', font = font1,  text='诊断图片', command=self.predict)
        #self.button_predict.pack(fill="x")
        self.button_predict.place(x=210, y=150)
        self.button_predict.config(state=tk.DISABLED)

        self.label_info = tk.Label(self.frame_left, font=('楷体', 14), justify=tk.LEFT, padx=2, pady=20)
        #self.label_info.pack(fill="x")
        self.label_info.place(x=400, y=250)
        self.label_info.config(text="请载入一个模型文件")
        self.canvas.create_text(360, 150, anchor='nw', font=('楷体', 13),fill = "white", text="step1:点击菜单栏，选择模型\n"
                                                                               "step2:点击加载模型按钮\n"
                                                                               "step3:点击上传图片按钮,选择png图片\n"
                                                                               "step4:点击诊断图片按钮\n")
        self.root.mainloop()

    def quit(self):
        self.destroy()

    def load_image_to_canvas(self, file_path):
        """把给定路径的图像加载入self.img 并绘制到canvas1"""
        def resize(w_box, h_box, pil_image):  # 参数是：要适应的窗口宽、高、Image.open后的图片
            w, h = pil_image.size  # 获取图像的原始大小
            f1 = 1.0 * w_box / w
            f2 = 1.0 * h_box / h
            factor = min([f1, f2])
            width = int(w * factor)
            height = int(h * factor)
            return pil_image.resize((width, height), Image.ANTIALIAS)
        try:
            img = Image.open(file_path)
            # print(img.size)
            self.img = img
            img_w, img_h = img.size
            img_w = 300
            img_h = 300
            img = resize(img_w, img_h, img)
            self.pil_img = ImageTk.PhotoImage(img)  # PhotoImage返回的对象必须一直被引用着，一旦失去引用，canvas上的图像立即消失
            self.canvas1.update()  # 获取宽高之前要先对于这个组件update()
            #x, y = 0, (self.canvas.winfo_height() - img.size[1]) / 2
            x, y = 0, 0
            self.canvas1.create_image(x, y, anchor='nw', image=self.pil_img)
        except Exception as e:
            self.label_info.config(text="图片载入出错")
        finally:
            self.button_predict.config(state=tk.NORMAL)
            self.label_info.config(text="图片已载入\n点击诊断按钮")

    def predict(self):
        """根据已载入的模型进行识别"""
        class_dict = {
            "0": "健康",
            "1": "非健康",
        }
        # 图片转换成224*224*3的格式
        preprocess_transform = transforms.Compose([
            # transforms.Resize(256),
            transforms.CenterCrop(224),
            transforms.ToTensor(),
            # 因为你的图像是单通道转成三通道
            transforms.Lambda(lambda x: x.repeat(3, 1, 1)),
            # transforms.ToPILImage(),
            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
        ])
        img = preprocess_transform(self.img)
        img.unsqueeze_(0)
        print(img.shape)

        with torch.no_grad():
            outputs, features = self.model(img)
            # outputs 对应属于每一类的概率值
            output = torch.squeeze(outputs)
            predict = torch.softmax(output, dim=0)
            predict_cla = torch.argmax(predict).numpy()  # 最大值位置索引
            print(predict_cla)
        class_str = class_dict[str(predict_cla)]
        prob_str = "%.1f" % (predict[predict_cla].item() * 100)
        self.label_info.config(text=f"类别：{class_str}\n可能性：{prob_str}%")

    def loadModel(self):
        """载入指定的模型"""
        try:

            global path
            path = self.modelPath
            if self.modelPath == "":
                return
            self.label_info.config(text="载入模型中……")
            self.model.load_state_dict(torch.load(self.modelPath))
            self.model.eval()
        except Exception as e:
            self.label_info.config(text="模型载入出错")
        finally:
            self.button_loadImage.config(state=tk.NORMAL)
            self.label_info.config(text="请打开一张图片")

    def loadModel1(self):
        """载入指定的模型"""
        self.modelPath = 'C:/Users/ASUS/Desktop/baseline (1)/model1/ResNet50/net_best.pkl'
        self.model = ResNet50(num_classes=2)
        self.button_loadModel.config(state=tk.NORMAL)
    def loadModel2(self):
        """载入指定的模型"""
        self.modelPath = 'C:/Users/ASUS/Desktop/baseline (1)/model1/VGG16_leaf/net_best.pkl'
        self.model =  VGG16()
        self.button_loadModel.config(state=tk.NORMAL)
    def loadModel3(self):
        """载入指定的模型"""
        self.modelPath = 'C:/Users/ASUS/Desktop/baseline (1)/model1/AlexNet_leaf/net_best.pkl'
        self.model = AlexNet()
        self.button_loadModel.config(state=tk.NORMAL)

    def loadimg(self):
        """载入指定的png图片"""
        default_dir = os.getcwd()
        photoPath = askopenfilename(title='打开一个照片（png格式）',
                                    initialdir=(os.path.expanduser(default_dir)),
                                    filetypes=[('png文件', '*.png'), ('All Files', '*')])
        if photoPath == "":
            return
        self.load_image_to_canvas(photoPath)


if __name__ == '__main__':
    win = MainWindow()
